/*
Two elements of a binary search tree (BST) are swapped by mistake.

Recover the tree without changing its structure.
*/

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */


/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    void recoverTree(TreeNode *root) {
        vector<TreeNode *> vals;
        int n1, n2;
        traverseInorder(root, vals);
        for (int i = 0; i < vals.size()-1; i++) {
            if (vals[i]->val > vals[i+1]->val) {
                n1 = i; break;
            }
        }
        for (int i = vals.size()-1; i > 0; i--) {
            if (vals[i]->val < vals[i-1]->val) {
                n2 = i; break;
            }
        }
        swap(vals[n1]->val, vals[n2]->val);
    }
private:
    void traverseInorder(TreeNode *n, vector<TreeNode *> &vals) {
        if (!n) return;
        traverseInorder(n->left, vals);
        vals.push_back(n);
        traverseInorder(n->right, vals);
    }
};

#if 0
class Solution {
public:
    void traverseInorder(TreeNode *n, vector<TreeNode *> &vals) {
        if (!n) return;
        traverseInorder(n->left, vals);
        vals.push_back(n);
        traverseInorder(n->right, vals);
    }
    void recoverTree(TreeNode *root) {
        /* O(n) solution */
        if (!root) return;
        vector<TreeNode *> vals;
        TreeNode *n1 = NULL;
        TreeNode *n2 = NULL;
        int n1_idx = 0;
        traverseInorder(root, vals);
        for (int i=0; i<vals.size()-1; i++) {
            if (vals[i]->val > vals[i+1]->val) {
                if (!n1) {
                    n1 = vals[i];
                    n1_idx = i;
                } else if (!n2) {
                    n2 = vals[i+1];
                }
            }
            if (n1 && n2) {break;}
        }
        // special case, two adjcent were swapped
        if (!n2) {n2 = vals[n1_idx+1];}
        // swap values of n1 and n2
        int temp = n1->val;
        n1->val = n2->val;
        n2->val = temp;
    }
};
#endif
